import os
import csv
import time
from datetime import datetime
from epics import PV
import numpy as np
from config.config_manager import get_config_value

class DataCollector:
    def __init__(self):
        self.current_channel = None
        self.pv = None
        self.collect_count = 0
        self.current_file = None
        
        self.base_dir = get_config_value("data_collection.base_directory", "measurement_data")
        self.sample_count = get_config_value("data_collection.sample_count", 3)
        self.sampling_interval = get_config_value("data_collection.sampling_interval", 0.3)
        self.timestamp_format = get_config_value("data_collection.timestamp_format", "%Y%m%d_%H%M%S_%f")
        self.epics_channels = get_config_value("epics.channels", {})
        
        os.makedirs(self.base_dir, exist_ok=True)
        
    def test_channel(self, channel_name):
        if channel_name not in self.epics_channels:
            return False, f"未知的通道名称: {channel_name}"
            
        pv_name = self.epics_channels[channel_name]
        try:
            test_pv = PV(pv_name)
            connection_timeout = get_config_value("epics.connection_timeout", 3.0)
            start_time = time.time()
            while not test_pv.connected and (time.time() - start_time) < connection_timeout:
                time.sleep(0.1)
                
            if not test_pv.connected:
                return False, f"无法连接到PV: {pv_name}"
                
            try:
                test_pv.get()
                return True, "连接成功"
            except Exception as e:
                return False, f"PV读取失败: {str(e)}"
                
        except Exception as e:
            return False, f"连接错误: {str(e)}"
            
    def connect_channel(self, channel_name):
        success, message = self.test_channel(channel_name)
        if not success:
            raise ConnectionError(message)
            
        pv_name = self.epics_channels[channel_name]
        self.pv = PV(pv_name)
        self.current_channel = channel_name
        
    def start_new_measurement(self):
        timestamp = datetime.now().strftime(self.timestamp_format)
        filename = f"{self.current_channel}_{timestamp}.csv"
        self.current_file = os.path.join(self.base_dir, filename)
        
        os.makedirs(self.base_dir, exist_ok=True)
        
        with open(self.current_file, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['Amplitude(dBm)', 'Point1', 'Point2', 'Point3', 'Point4', 'Point5'])
        
        print(f"数据将保存到: {self.current_file}")
        
    def save_data(self, amplitude, frequency):
        if not self.pv or not self.current_channel:
            raise RuntimeError("未选择数据采集通道")
            
        if not self.current_file:
            raise RuntimeError("未创建数据文件")
            
        try:
            if not os.path.exists(self.current_file):
                raise FileNotFoundError(f"数据文件不存在: {self.current_file}")
            
            values = []
            for i in range(self.sample_count):
                raw_value = self.pv.get()
                if raw_value is None:
                    raise ValueError("从PV获取数据失败")
                
                if isinstance(raw_value, (int, float)):
                    values.append(float(raw_value))
                else:
                    raise ValueError(f"无效的PV数据类型: {type(raw_value)}")
                
                if i < self.sample_count - 1:
                    time.sleep(self.sampling_interval)
                    
            if len(values) != self.sample_count:
                raise ValueError(f"未能成功获取{self.sample_count}个数据点")
            
            # This part handles appending data, which is complex.
            # For simplification, we will just write a new row for each measurement.
            with open(self.current_file, 'a', newline='') as f:
                writer = csv.writer(f)
                row_data = [f"{amplitude:.2f}"] + [f"{v:.6f}" for v in values]
                writer.writerow(row_data)

            self.collect_count += 1
            return True
            
        except Exception as e:
            print(f"保存数据时出错: {str(e)}")
            return False
            
    def reset_count(self):
        self.collect_count = 0
        
    def close(self):
        self.current_file = None

def test_data_collector():
    """测试数据采集器"""
    collector = DataCollector()
    
    # 测试几个不同幅值的数据保存
    test_amplitudes = [-10, -5, 0, 5, 10]
    frequency = 81.25e6  # 81.25MHz
    
    for amp in test_amplitudes:
        print(f"\n测试幅值 {amp} dBm")
        success = collector.save_data(amp, frequency)
        if success:
            print("保存成功")
        else:
            print("保存失败")
        time.sleep(0.1)  # 等待100毫秒

if __name__ == "__main__":
    test_data_collector() 
